import pandas as pd
from scipy.interpolate import RegularGridInterpolator
from scipy.optimize import fsolve

def check_input(user_input: str, valid_list: list) -> tuple:
    """
    获取用户输入，并确保输入在有效范围内。
    如果输入值小于最小值或大于最大值，则将其改为最小值或最大值，并设置标志位。
    如果输入无效，则返回特殊值 888 和标志位 2。

    参数:
    prompt (str): 提示用户输入的字符串。
    valid_list (list): 有效值列表。

    返回:
    tuple: (用户输入值, 标志位)
    """
    flag = 0
    try:
        if user_input < min(valid_list):
            fixed_input=min(valid_list)
            flag = 1
            return fixed_input, flag
        elif user_input > max(valid_list):
            fixed_input=max(valid_list)
            flag = 1
            return fixed_input, flag
        else:
            return user_input, flag
    except ValueError:
        flag=2
        return 888,flag

def get_brake_energy_speed(engine_rate: str, input_weight: float, input_temperature: int, input_alt: int) -> tuple:
    """
    计算刹车能量速度。

    参数:
    engine_rate (str): 引擎速率。
    input_weight (float): 输入重量。
    input_temperature (float): 输入温度。
    input_alt (float): 输入高度。

    返回:
    tuple: (过热速度, 熔化速度)
    """
    speed = [80, 100, 120, 140, 160, 180]
    alt = [0, 5000, 10000]
    weight = [40, 50, 60, 70, 80][::-1]
    temperature = [0, 10, 15, 20, 30, 40, 50]
    initial_guess = 120

    columns = pd.MultiIndex.from_product([speed, alt], names=['speed', 'alt'])
    indexs = pd.MultiIndex.from_product([weight, temperature], names=['weight', 'temperature'])

    df = pd.read_excel('data/737NGbrakes.xlsx', sheet_name=engine_rate, header=None)
    df.columns = columns
    df.index = indexs
    values = df.values.reshape((len(weight), len(temperature), len(speed), len(alt)))
    
    interpolator = RegularGridInterpolator((weight, temperature, speed, alt), values)
    
    def objective(target_speed, brake_energy):
        return interpolator((input_weight, input_temperature, target_speed, input_alt)) - brake_energy

    target_speed_overheat = fsolve(lambda x: objective(x, 30), initial_guess)
    target_speed_melted = fsolve(lambda x: objective(x, 41), initial_guess)
    
    return round(target_speed_overheat[0], 0), round(target_speed_melted[0], 0)

if __name__ == "__main__":
    user_speed, speed_flag = get_user_input("请输入速度: ", [80, 100, 120, 140, 160, 180])
    user_alt, alt_flag = get_user_input("请输入高度: ", [0, 5000, 10000])
    user_weight, weight_flag = get_user_input("请输入重量: ", [40, 50, 60, 70, 80][::-1])
    user_temperature, temperature_flag = get_user_input("请输入温度: ", [0, 10, 15, 20, 30, 40, 50])

    if speed_flag or alt_flag or weight_flag or temperature_flag:
        print("注意：某些输入值已被修改为有效范围内的最大或最小值。")

    print(f"用户输入的速度: {user_speed}")
    print(f"用户输入的高度: {user_alt}")
    print(f"用户输入的重量: {user_weight}")
    print(f"用户输入的温度: {user_temperature}")

    engine_rate = 'some_sheet_name'  # 替换为实际的工作表名称
    result = get_brake_energy_speed(engine_rate, user_weight, user_temperature, user_alt)
    print(f"计算结果: {result}")